import _thread as thread
import base64
import datetime
import hashlib
import hmac
import json
from datetime import datetime
from time import mktime
from urllib.parse import urlencode
from wsgiref.handlers import format_date_time
import websocket

class SparkAPI:
    def __init__(self, app_id, api_key, api_secret):
        self.APPID = app_id
        self.APIKey = api_key
        self.APISecret = api_secret
        self.host = "maas-api.cn-huabei-1.xf-yun.com"
        self.path = "/v1.1/chat"

    def _create_url(self):
        # 生成RFC1123格式的时间戳
        now = datetime.now()
        date = format_date_time(mktime(now.timetuple()))

        # 拼接字符串
        signature_origin = f"host: {self.host}\ndate: {date}\nGET {self.path} HTTP/1.1"

        # 进行hmac-sha256进行加密
        signature_sha = hmac.new(self.APISecret.encode('utf-8'),
                                signature_origin.encode('utf-8'),
                                digestmod=hashlib.sha256).digest()

        signature_sha_base64 = base64.b64encode(signature_sha).decode(encoding='utf-8')

        authorization_origin = f'api_key="{self.APIKey}", algorithm="hmac-sha256", headers="host date request-line", signature="{signature_sha_base64}"'

        authorization = base64.b64encode(authorization_origin.encode('utf-8')).decode(encoding='utf-8')

        # 将请求的鉴权参数组合为字典
        v = {
            "authorization": authorization,
            "date": date,
            "host": self.host
        }
        # 拼接鉴权参数，生成url
        url = f"wss://{self.host}{self.path}?{urlencode(v)}"
        return url

    def _create_message(self, question):
        data = {
            "header": {
                "app_id": self.APPID,
                "patch_id": ["0"]
            },
            "parameter": {
                "chat": {
                    "domain": "xdeepseekr1",
                    "temperature": 0.5,
                    "top_k": 4
                }
            },
            "payload": {
                "message": {
                    "text": [
                        {"role": "user", "content": question}
                    ]
                }
            }
        }
        return json.dumps(data)

    def chat(self, question, on_message=None, on_error=None):
        url = self._create_url()
        ws = websocket.WebSocketApp(
            url,
            on_message=on_message if on_message else self._on_message,
            on_error=on_error if on_error else self._on_error
        )
        ws.on_open = lambda ws: self._on_open(ws, question)
        ws.run_forever()

    def _on_message(self, ws, message):
        data = json.loads(message)
        code = data['header']['code']
        if code != 0:
            error_msg = data["header"]["message"]
            print(f'请求错误: {code}, {data}')
            ws.close()
            return

        choices = data["payload"]["choices"]
        status = choices["status"]
        content = choices["text"][0]["content"]
        print(content, end='')

        # 最后一条消息
        if status == 2:
            ws.close()

    def _on_error(self, ws, error):
        print(f"### error: {error}")

    def _on_open(self, ws, question):
        message = self._create_message(question)
        def run():
            ws.send(message)
        thread.start_new_thread(run, ())

# 使用示例
def main():
    # 请替换为您的API密钥信息
    app_id = "xx"
    api_key = "xx"
    api_secret = "xx"

    api = SparkAPI(app_id, api_key, api_secret)
    api.chat("使用3种方法证明毕达哥拉斯定理")

if __name__ == "__main__":
    main()